import os
import cv2
from skimage import measure, filters, io
from skimage.filters import rank
import matplotlib.pyplot as plt
from skimage.morphology import disk
from utils import DCT2, IDCT2, blkproc, wavelet_haar, inv_wavelet_haar, PCAdenoising
import numpy as np
from sklearn.decomposition import PCA


"""
1. compare the noised image and the source img
"""
ORIGINAL = 'original'
GAUSSIAN = 'gaussian'
SALT_PEPPER = 'salt-pepper'
SPECKLE = 'speckle'
#  gaussian
for i in os.listdir(ORIGINAL):
    img_path = GAUSSIAN + '/' + i[0:-4] + '_' + GAUSSIAN + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    print 'PSNR-' + i + ':  ' + str(measure.compare_psnr(im_true=img_arr_org, im_test=img_arr, dynamic_range=255))
print 'gaussian'
print '*' * 50

#  salt-pepper
for i in os.listdir(ORIGINAL):
    img_path = SALT_PEPPER + '/' + i[0:-4] + '_' + SALT_PEPPER + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    print 'PSNR-' + i + ':  ' + str(measure.compare_psnr(im_true=img_arr_org, im_test=img_arr, dynamic_range=255))
print 'salt-pepper'
print '*' * 50

#  speckle
for i in os.listdir(ORIGINAL):
    img_path = SPECKLE + '/' + i[0:-4] + '_' + SPECKLE + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    print 'PSNR-' + i + ':  ' + str(measure.compare_psnr(im_true=img_arr_org, im_test=img_arr, dynamic_range=255))
print 'speckle'
print '*' * 50

"""
2. median_filter and mean_filter
"""
#  gaussian
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = GAUSSIAN + '/' + i[0:-4] + '_' + GAUSSIAN + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_median = filters.median(img_arr, disk(3))
    img_arr_mean = rank.mean(img_arr, disk(3))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_median)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_mean)
    plot_flag += 1
    print 'PSNR(median_mean)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_median, dynamic_range=255))
    print 'PSNR(mean)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_mean, dynamic_range=255))
    print '-' * 30
fig1.savefig('res/median_mean/gaussian.png')
print 'GAUSSIAN - filter'
print '*' * 50

#  salt-pepper
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SALT_PEPPER + '/' + i[0:-4] + '_' + SALT_PEPPER + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_median = filters.median(img_arr, disk(3))
    img_arr_mean = rank.mean(img_arr, disk(3))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_median)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_mean)
    plot_flag += 1
    print 'PSNR(median_mean)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_median, dynamic_range=255))
    print 'PSNR(mean)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_mean, dynamic_range=255))
    print '-' * 30
fig1.savefig('res/median_mean/salt_pepper.png')
print 'SALT_PEPPER - filter'
print '*' * 50

#  speckle
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SPECKLE + '/' + i[0:-4] + '_' + SPECKLE + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_median = filters.median(img_arr, disk(3))
    img_arr_mean = rank.mean(img_arr, disk(3))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_median)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_mean)
    plot_flag += 1
    print 'PSNR(median_mean)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_median, dynamic_range=255))
    print 'PSNR(mean)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_mean, dynamic_range=255))
    print '-' * 30
fig1.savefig('res/median_mean/speckle.png')
print 'SPECKLE - filter'
print '*' * 50

"""
3. DCT
"""
#  gaussian
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = GAUSSIAN + '/' + i[0:-4] + '_' + GAUSSIAN + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_dct = np.uint8(blkproc(img_arr, (8, 8), DCT2))
    img_arr_idct = np.uint8(blkproc(img_arr, (8, 8), IDCT2))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_dct)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_idct)
    plot_flag += 1
    print 'PSNR(DCT)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_idct, dynamic_range=255))

fig1.savefig('res/dct/gaussian.png')
print 'GAUSSIAN - dct'
print '*' * 50

#  salt-pepper
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SALT_PEPPER + '/' + i[0:-4] + '_' + SALT_PEPPER + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_dct = np.uint8(blkproc(img_arr, (8, 8), DCT2))
    img_arr_idct = np.uint8(blkproc(img_arr, (8, 8), IDCT2))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_dct)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_idct)
    plot_flag += 1
    print 'PSNR(DCT)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_idct, dynamic_range=255))

fig1.savefig('res/dct/SALT_PEPPER.png')
print 'SALT_PEPPER - dct'
print '*' * 50

#  SPECKLE
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SPECKLE + '/' + i[0:-4] + '_' + SPECKLE + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_dct = np.uint8(blkproc(img_arr, (8, 8), DCT2))
    img_arr_idct = np.uint8(blkproc(img_arr, (8, 8), IDCT2))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_dct)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_idct)
    plot_flag += 1
    print 'PSNR(DCT)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_idct, dynamic_range=255))

fig1.savefig('res/dct/SPECKLE.png')
print 'SPECKLE - dct'
print '*' * 50

"""
4. haar wavelet
"""
#  gaussian
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = GAUSSIAN + '/' + i[0:-4] + '_' + GAUSSIAN + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_dct = np.uint8(wavelet_haar(img_arr, 100))
    img_arr_idct = np.uint8(inv_wavelet_haar(img_arr, 100))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_dct)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_idct)
    plot_flag += 1
    print 'PSNR(WLT)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_idct, dynamic_range=255))

fig1.savefig('res/wlt/gaussian.png')
print 'GAUSSIAN - WLT'
print '*' * 50

#  salt-pepper
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SALT_PEPPER + '/' + i[0:-4] + '_' + SALT_PEPPER + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_dct = np.uint8(wavelet_haar(img_arr, 100))
    img_arr_idct = np.uint8(inv_wavelet_haar(img_arr, 100))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_dct)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_idct)
    plot_flag += 1
    print 'PSNR(WLT)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_idct, dynamic_range=255))

fig1.savefig('res/wlt/SALT_PEPPER.png')
print 'SALT_PEPPER - WLT'
print '*' * 50

#  SPECKLE
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SPECKLE + '/' + i[0:-4] + '_' + SPECKLE + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_dct = np.uint8(wavelet_haar(img_arr, 100))
    img_arr_idct = np.uint8(inv_wavelet_haar(img_arr, 100))
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_dct)
    plot_flag += 1
    fig1.add_subplot(3, 4, plot_flag)
    io.imshow(img_arr_idct)
    plot_flag += 1
    print 'PSNR(WLT)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_idct, dynamic_range=255))

fig1.savefig('res/wlt/SPECKLE.png')
print 'SPECKLE - WLT'
print '*' * 50

"""
5. PCA
"""

# gaussian
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = GAUSSIAN + '/' + i[0:-4] + '_' + GAUSSIAN + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_pca = np.uint8(PCAdenoising(img_arr, (8, 8), 6))
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr_pca)
    plot_flag += 1
    print 'PSNR(PCA)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_pca, dynamic_range=255))

fig1.savefig('res/pca/gaussian.png')
print 'GAUSSIAN - PCA'
print '*' * 50

# SALT_PEPPER
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SALT_PEPPER + '/' + i[0:-4] + '_' + SALT_PEPPER + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_pca = np.uint8(PCAdenoising(img_arr, (8, 8), 6))
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr_pca)
    plot_flag += 1
    print 'PSNR(PCA)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_pca, dynamic_range=255))

fig1.savefig('res/pca/SALT_PEPPER.png')
print 'SALT_PEPPER - PCA'
print '*' * 50

# SPECKLE
fig1 = plt.figure()
plot_flag = 1
for i in os.listdir(ORIGINAL):
    img_path = SPECKLE + '/' + i[0:-4] + '_' + SPECKLE + '.png'
    img_path_org = ORIGINAL + '/' + i
    img_arr = cv2.imread(img_path, 0)
    img_arr_org = cv2.imread(img_path_org, 0)
    img_arr_pca = np.uint8(PCAdenoising(img_arr, (8, 8), 6))
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr)
    plot_flag += 1
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr_org)
    plot_flag += 1
    fig1.add_subplot(3, 3, plot_flag)
    io.imshow(img_arr_pca)
    plot_flag += 1
    print 'PSNR(PCA)-' + i + ':  ' + str(
        measure.compare_psnr(im_true=img_arr_org, im_test=img_arr_pca, dynamic_range=255))

fig1.savefig('res/pca/SPECKLE.png')
print 'SPECKLE - PCA'
print '*' * 50